#if __CUDA_ARCH__ == 100
# error Atomics only used with > sm_10 architecture
#endif
#include "cuPrintf.h"

// This is the smallest amount of memory, per-thread, which is allowed. It is also the largest amount of space a single printf() can take up
const static int CUPRINTF_MAX_LEN = 256;

// This structure is used internally to track block/thread output restrictions.
typedef struct __align__(8) {
	int threadid;				// CUPRINTF_UNRESTRICTED for unrestricted
	int blockid;				// CUPRINTF_UNRESTRICTED for unrestricted
} cuPrintfRestriction;

// This is the header preceeding all printf entries. NOTE: It *must* be size-aligned to the maximum entity size (size_t)
typedef struct __align__(8) {
	unsigned short magic;                   // Magic number says we're valid
	unsigned short fmtoffset;               // Offset of fmt string into buffer
	unsigned short blockid;                 // Block ID of author
	unsigned short threadid;                // Thread ID of author
} cuPrintfHeader;

// Special header for sm_10 architecture
#define CUPRINTF_SM10_MAGIC   0xC810        // Not a valid ascii character
typedef struct __align__(16) {
	unsigned short magic;                   // sm_10 specific magic number
	unsigned short unused;
	unsigned int thread_index;              // thread ID for this buffer
	unsigned int thread_buf_len;            // per-thread buffer length
	unsigned int offset;                    // most recent printf's offset
} cuPrintfHeaderSM10;

typedef struct __align__(8) {
	
} cuPrintfCtx;

// Because we can't write an element which is not aligned to its bit-size, we have to align all sizes and variables on maximum-size boundaries.
// That means sizeof(double) in this case, but we'll use (long long) for better arch<1.3 support
#define CUPRINTF_ALIGN_SIZE      sizeof(long long)

// All our headers are prefixed with a magic number so we know they're ready
#define CUPRINTF_SM11_MAGIC  (unsigned short)0xC811        // Not a valid ascii character

///////////////////////////////////////////////////////////////////////////////
// HOST SIDE

#include <stdio.h>
static FILE *printf_fp;
static char *printfbuf_start = NULL;
static char *printfbuf_device = NULL;
static int printfbuf_len = 0;

//
//  outputPrintfData
//  Our own internal function, which takes a pointer to a data buffer and passes it through libc's printf for output.
//  We receive the formate string and a pointer to where the data is held. We then run through and print it out.
//  Returns 0 on failure, 1 on success
//
static int outputPrintfData(char *fmt, char *data)
{
	// Format string is prefixed by a length that we don't need
	fmt += CUPRINTF_ALIGN_SIZE;

	// Now run through it, printing everything we can. We must run to every % character, extract only that, and use printf to format it.
	char *p = strchr(fmt, '%');
	while (p != NULL)
	{
		// Print up to the % character
		*p = '\0'; fputs(fmt, printf_fp); *p = '%'; // Put back the %

		// Now handle the format specifier
		char *format = p++;         // Points to the '%'
		p += strcspn(p, "%cdiouxXeEfgGaAnps");
		if (*p == '\0') { fmt = format; break; } // If no format specifier, print the whole thing

		// Cut out the format bit and use printf to print it. It's prefixed by its length.
		int arglen = *(int *)data;
		if (arglen > CUPRINTF_MAX_LEN) { fputs("Corrupt printf buffer data - aborting\n", printf_fp); return 0; }

		data += CUPRINTF_ALIGN_SIZE;

		char specifier = *p++;
		char c = *p;        // Store for later
		*p = '\0';
		switch(specifier)
		{
		// These all take integer arguments
		case 'c': case 'd': case 'i': case 'o': case 'u': case 'x': case 'X': case 'p': fprintf(printf_fp, format, *((int *)data)); break;

		// These all take double arguments
		case 'e': case 'E': case 'f': case 'g': case 'G': case 'a': case 'A':
			// Float vs. Double thing
			if (arglen == 4) fprintf(printf_fp, format, *((float *)data)); else fprintf(printf_fp, format, *((double *)data)); break;

		// Strings are handled in a special way
		case 's': fprintf(printf_fp, format, (char *)data); break;

		// % is special
		case '%': fprintf(printf_fp, "%%"); break;

		// Everything else is just printed out as-is
		default: fprintf(printf_fp, format); break;
		}
		data += CUPRINTF_ALIGN_SIZE;    // Move on to next argument
		*p = c;                     // Restore what we removed
		fmt = p;                    // Adjust fmt string to be past the specifier
		p = strchr(fmt, '%');       // and get the next specifier
	}

	// Print out the last of the string
	fputs(fmt, printf_fp);
	return 1;
}


//
//  doPrintfDisplay
//
//  This runs through the blocks of CUPRINTF_MAX_LEN-sized data, calling the print function above to display them. We've got this separate from
//  cudaPrintfDisplay() below so we can handle the SM_10 architecture partitioning.
//
static int doPrintfDisplay(int headings, int clear, char *bufstart, char *bufend, char *bufptr, char *endptr)
{
	// Grab, piece-by-piece, each output element until we catch up with the circular buffer end pointer
	int printf_count = 0;
	char printfbuf_local[CUPRINTF_MAX_LEN+1];
	printfbuf_local[CUPRINTF_MAX_LEN] = '\0';

	while (bufptr != endptr)
	{
		// Wrap ourselves at the end-of-buffer
		if (bufptr == bufend)
			bufptr = bufstart;

		// Adjust our start pointer to within the circular buffer and copy a block.
		cudaMemcpy(printfbuf_local, bufptr, CUPRINTF_MAX_LEN, cudaMemcpyDeviceToHost);

		// If the magic number isn't valid, then this write hasn't gone through yet and we'll wait until it does (or we're past the end for non-async printfs).
		cuPrintfHeader *hdr = (cuPrintfHeader *)printfbuf_local;
		if (hdr->magic != CUPRINTF_SM11_MAGIC || hdr->fmtoffset >= CUPRINTF_MAX_LEN)
		{
			//fprintf(printf_fp, "Bad magic number in printf header\n");
			break;
		}

		// Extract all the info and get this printf done
		if (headings) fprintf(printf_fp, "[%d, %d]: ", hdr->blockid, hdr->threadid);
		if (hdr->fmtoffset == 0) fprintf(printf_fp, "printf buffer overflow\n");
		else if (!outputPrintfData(printfbuf_local + hdr->fmtoffset, printfbuf_local + sizeof(cuPrintfHeader))) break;
		printf_count++;

		// Clear if asked
		if (clear) cudaMemset(bufptr, 0, CUPRINTF_MAX_LEN);

		// Now advance our start location, because we're done, and keep copying
		bufptr += CUPRINTF_MAX_LEN;
	}

	return printf_count;
}


//
//  cudaPrintfInit
//
//  Takes a buffer length to allocate, creates the memory on the device and
//  returns a pointer to it for when a kernel is called. It's up to the caller
//  to free it.
//
extern "C" cudaError_t cudaPrintfInit(size_t bufferLen)
{
	cudaError_t cudaStatus;
	// Fix up bufferlen to be a multiple of CUPRINTF_MAX_LEN
	bufferLen = (bufferLen < CUPRINTF_MAX_LEN ? CUPRINTF_MAX_LEN : bufferLen);
	if ((bufferLen % CUPRINTF_MAX_LEN) > 0)
		bufferLen += (CUPRINTF_MAX_LEN - (bufferLen % CUPRINTF_MAX_LEN));
	printfbuf_len = (int)bufferLen;

	// Allocate a print buffer on the device and zero it
	if (cudaMalloc((void **)&printfbuf_device, printfbuf_len) != cudaSuccess) return cudaErrorInitializationError;
	cudaMemset(printfbuf_device, 0, printfbuf_len);
	printfbuf_start = printfbuf_device; // Where we start reading from

	// No restrictions to begin with
	cuPrintfRestriction restrict;
	restrict.threadid = restrict.blockid = CUPRINTF_UNRESTRICTED;


	// Allocate a print buffer on the device and zero it
	fallocHeap* heap;
	if ((!error && (cudaMalloc((void**)&heap, length) != cudaSuccess)) ||
		(error && ((*error = cudaMalloc((void**)&heap, length)) != cudaSuccess)))
		return host;
	cudaMemset(heap, 0, length);
	// transfer to heap
	fallocHeap hostHeap;
	hostHeap.chunks = chunks;
	hostHeap.freeChunks = nullptr;
	hostHeap.reserved = reserved;
	if ((!error && (cudaMemcpy(heap, &hostHeap, sizeof(fallocHeap), cudaMemcpyHostToDevice) != cudaSuccess)) ||
		(error && ((*error = cudaMemcpy(heap, &hostHeap, sizeof(fallocHeap), cudaMemcpyHostToDevice)) != cudaSuccess)))
		return host;
	// return the heap
	if (error)
		*error = cudaSuccess;
	host.heap = heap;
	host.length = (int)length;
	host.reserved = reserved;
	return host;
}


	//if ((cudaStatus = cudaMemcpyToSymbol(restrictRules, &restrict, sizeof(restrict))) != cudaSuccess) return cudaStatus;

	//// Initialise the buffer and the respective lengths/pointers.
	//if ((cudaStatus = cudaMemcpyToSymbol(globalPrintfBuffer, &printfbuf_device, sizeof(char *))) != cudaSuccess) return cudaStatus;
	//if ((cudaStatus = cudaMemcpyToSymbol(printfBufferPtr, &printfbuf_device, sizeof(char *))) != cudaSuccess) return cudaStatus;
	//if ((cudaStatus = cudaMemcpyToSymbol(printfBufferLength, &printfbuf_len, sizeof(printfbuf_len))) != cudaSuccess) return cudaStatus;

	return cudaSuccess;
}


//
//  cudaPrintfEnd
//
//  Frees up the memory which we allocated
//
extern "C" void cudaPrintfEnd()
{
	if (!printfbuf_start || !printfbuf_device) return;
	cudaFree(printfbuf_device);
	printfbuf_start = printfbuf_device = NULL;
}


//
//  cudaPrintfDisplay
//
//  Each call to this function dumps the entire current contents of the printf buffer to the pre-specified FILE pointer. The
//	circular "start" pointer is advanced so that subsequent calls dumps only new stuff.
//
//  In the case of async memory access (via streams), call this repeatedly to keep trying to empty the buffer. If it's a sync
//  access, then the whole buffer should empty in one go.
//
//	Arguments:
//		outputFP     - File descriptor to output to (NULL => stdout)
//		showThreadID - If true, prints [block,thread] before each line
//
extern "C" cudaError_t cudaPrintfDisplay(void *outputFP, bool showThreadID)
{
	printf_fp = (FILE *)(outputFP == NULL ? stdout : outputFP);

	// For now, we force "synchronous" mode which means we're not concurrent with kernel execution. This also means we don't need clearOnPrint.
	// If you're patching it for async operation, here's where you want it.
	bool sync_printfs = true;
	bool clearOnPrint = false;

	// Initialisation check
	if (!printfbuf_start || !printfbuf_device || !printf_fp) return cudaErrorMissingConfiguration;

	// To determine which architecture we're using, we read the first short from the buffer - it'll be the magic number relating to the version.
	unsigned short magic;
	cudaMemcpy(&magic, printfbuf_device, sizeof(unsigned short), cudaMemcpyDeviceToHost);

	// For SM_10 architecture, we've split our buffer into one-per-thread. That means we must do each thread block separately. It'll require
	// extra reading. We also, for now, don't support async printfs because that requires tracking one start pointer per thread.
	if (magic == CUPRINTF_SM10_MAGIC)
	{
		sync_printfs = true;
		clearOnPrint = false;
		int blocklen = 0;
		char *blockptr = printfbuf_device;
		while(blockptr < (printfbuf_device + printfbuf_len))
		{
			cuPrintfHeaderSM10 hdr;
			cudaMemcpy(&hdr, blockptr, sizeof(hdr), cudaMemcpyDeviceToHost);

			// We get our block-size-step from the very first header
			if(hdr.thread_buf_len != 0)
				blocklen = hdr.thread_buf_len;

			// No magic number means no printfs from this thread
			if(hdr.magic != CUPRINTF_SM10_MAGIC)
			{
				if(blocklen == 0)
				{
					fprintf(printf_fp, "No printf headers found at all!\n");
					break;                              // No valid headers!
				}
				blockptr += blocklen;
				continue;
			}

			// "offset" is non-zero then we can print the block contents
			if(hdr.offset > 0)
			{
				// For synchronous printfs, we must print from endptr->bufend, then from start->end
				if(sync_printfs)
					doPrintfDisplay(showThreadID, clearOnPrint, blockptr+CUPRINTF_MAX_LEN, blockptr+hdr.thread_buf_len, blockptr+hdr.offset+CUPRINTF_MAX_LEN, blockptr+hdr.thread_buf_len);
				doPrintfDisplay(showThreadID, clearOnPrint, blockptr+CUPRINTF_MAX_LEN, blockptr+hdr.thread_buf_len, blockptr+CUPRINTF_MAX_LEN, blockptr+hdr.offset+CUPRINTF_MAX_LEN);
			}

			// Move on to the next block and loop again
			blockptr += hdr.thread_buf_len;
		}
	}
	// For SM_11 and up, everything is a single buffer and it's simple
	else if(magic == CUPRINTF_SM11_MAGIC)
	{
		// Grab the current "end of circular buffer" pointer.
		char *printfbuf_end = NULL;
		cudaMemcpyFromSymbol(&printfbuf_end, "printfBufferPtr", sizeof(char *));

		// Adjust our starting and ending pointers to within the block
		char *bufptr = ((printfbuf_start - printfbuf_device) % printfbuf_len) + printfbuf_device;
		char *endptr = ((printfbuf_end - printfbuf_device) % printfbuf_len) + printfbuf_device;

		// For synchronous (i.e. after-kernel-exit) printf display, we have to handle circular
		// buffer wrap carefully because we could miss those past "end".
		if(sync_printfs)
			doPrintfDisplay(showThreadID, clearOnPrint, printfbuf_device, printfbuf_device+printfbuf_len, endptr, printfbuf_device+printfbuf_len);
		doPrintfDisplay(showThreadID, clearOnPrint, printfbuf_device, printfbuf_device+printfbuf_len, bufptr, endptr);

		printfbuf_start = printfbuf_end;
	}
	else
		;//printf("Bad magic number in cuPrintf buffer header\n");

	// If we were synchronous, then we must ensure that the memory is cleared on exit
	// otherwise another kernel launch with a different grid size could conflict.
	if(sync_printfs)
		cudaMemset(printfbuf_device, 0, printfbuf_len);

	return cudaSuccess;
}
